Step FunctionsでDynamoDBのエクスポート機能を活用して、Athenaでクエリする環境作成を自動化してみた

    const articeTable = new dynamodb.Table(this, `Article`, {
      partitionKey: {
        name: 'userId',
        type: dynamodb.AttributeType.STRING,
      },
      sortKey: {
        name: 'articleId',
        type: dynamodb.AttributeType.STRING,
      },
      tableName: `Article`,
      pointInTimeRecovery: true, // エクスポート機能を使う場合は、有効にする必要があります。
      billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
    });

{
  "userId": "7155dc64-982a-4559-9f94-5320dc4d003d",
  "articleId": "XXXXXXX",
  "content": "hoge",
  "createAt": 1577804400
}

{
  "userId": "7155dc64-982a-4559-9f94-5320dc4d003d",
  "articleId": "YYYYYY",
  "content": "foo",
  "createAt": 1577804401
}

cdk from '@aws-cdk/core';
import * as s3 from '@aws-cdk/aws-s3';
import * as lambda from '@aws-cdk/aws-lambda';
import * as iam from '@aws-cdk/aws-iam';
import * as events from '@aws-cdk/aws-events';
import * as eventsTargets from '@aws-cdk/aws-events-targets';
import * as athena from '@aws-cdk/aws-athena';
import * as glue from '@aws-cdk/aws-glue';
import * as stepfunctions from '@aws-cdk/aws-stepfunctions';
import * as stepfunctionsTasks from '@aws-cdk/aws-stepfunctions-tasks';

export class AnalysisStack extends cdk.Stack {
  constructor(scope: cdk.Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const accountId = cdk.Stack.of(this).account;
    const DB_NAME = `ddb`;

    // DynamoDB出力用、Athena検索用のバケットを作成
    const anlysisBucket = new s3.Bucket(this, 'AnalysisBucket', {
      bucketName: `analysis-bucket-${accountId}`,
    });

    // Athenaのワークグループを作成、Lambdaからクエリを実行する際に必要
    const athenWorkGroup = new athena.CfnWorkGroup(this, `AthenaWorkGroup`, {
      name: `work-group`,
      workGroupConfiguration: {
        resultConfiguration: {
          outputLocation: `s3://${anlysisBucket.bucketName}/athena-work-group/`,
        },
      },
    });

    // Lambda①
    const exportDynamoFunction = new lambda.Function(
      this,
      'exportDynamoFunction',
      {
        code: lambda.Code.fromAsset(`dist/exportDynamo`),
        environment: {
          EXPORT_S3_BUCKET_NAME: anlysisBucket.bucketName,
        },
        functionName: `export-article-table`,
        handler: 'index.handler',
        runtime: lambda.Runtime.NODEJS_14_X,
        tracing: lambda.Tracing.ACTIVE,
        timeout: cdk.Duration.seconds(10),
      },
    );
    exportDynamoFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['dynamodb:exportTableToPointInTime'],
        resources: ['*'],
      }),
    );
    exportDynamoFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['s3:PutObject'],
        resources: ['*'],
      }),
    );

    // Lambda②
    const getExportDynamoStatusFunction = new lambda.Function(
      this,
      'getExportDynamoStatusFunction',
      {
        code: lambda.Code.fromAsset(`dist/getExportDynamoStatus`),
        environment: {},
        functionName: `get-export-dynamo-status`,
        handler: 'index.handler',
        runtime: lambda.Runtime.NODEJS_14_X,
        tracing: lambda.Tracing.ACTIVE,
        timeout: cdk.Duration.seconds(10),
      },
    );
    getExportDynamoStatusFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['dynamodb:*'],
        resources: ['*'],
      }),
    );

    // Lambda③
    const relocateTableFunction = new lambda.Function(
      this,
      'relocateTableFunction',
      {
        code: lambda.Code.fromAsset(`dist/relocateTable`),
        environment: {
          DB_NAME: DB_NAME,
          ATHENA_WORK_GROUP: athenWorkGroup.name,
        },
        functionName: `relocate-table`,
        handler: 'index.handler',
        runtime: lambda.Runtime.NODEJS_14_X,
        tracing: lambda.Tracing.ACTIVE,
        timeout: cdk.Duration.seconds(10),
      },
    );
    relocateTableFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['athena:*', 's3:*', 'glue:*'],
        resources: ['*'],
      }),
    );

    // Lambda④
    const getRelocateTableStatusFunction = new lambda.Function(
      this,
      'getRelocateTableStatusFunction',
      {
        code: lambda.Code.fromAsset(`dist/getRelocateTableStatus`),
        functionName: `relocate-table-status`,
        handler: 'index.handler',
        runtime: lambda.Runtime.NODEJS_14_X,
        tracing: lambda.Tracing.ACTIVE,
        timeout: cdk.Duration.seconds(10),
      },
    );
    getRelocateTableStatusFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['athena:*', 's3:*'],
        resources: ['*'],
      }),
    );

    // ステートマシーンの作成
    const EXPORT_DYNAMO_WAIT_TIME = 5;
    const RELOCATE_TABLE_WAIT_TIME = 1;

    const exportDynamoTask = new stepfunctionsTasks.LambdaInvoke(
      this,
      'Request export of DynamoDB to S3',
      {
        lambdaFunction: exportDynamoFunction,
        invocationType:
          stepfunctionsTasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload',
      },
    );

    const exportDynamoWait = new stepfunctions.Wait(
      this,
      'Wait to request export of DynamoDB to S3',
      {
        time: stepfunctions.WaitTime.duration(
          cdk.Duration.minutes(EXPORT_DYNAMO_WAIT_TIME),
        ),
      },
    );

    const getExportDynamoStatusTask = new stepfunctionsTasks.LambdaInvoke(
      this,
      'Get export of DynamoDB to S3 status',
      {
        lambdaFunction: getExportDynamoStatusFunction,
        invocationType:
          stepfunctionsTasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload',
      },
    );

    const relocateTableTask = new stepfunctionsTasks.LambdaInvoke(
      this,
      'Query to change location',
      {
        lambdaFunction: relocateTableFunction,
        invocationType:
          stepfunctionsTasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload',
      },
    );

    const getRelocateTableTask = new stepfunctionsTasks.LambdaInvoke(
      this,
      'Get query status',
      {
        lambdaFunction: getRelocateTableStatusFunction,
        invocationType:
          stepfunctionsTasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload',
      },
    );

    const relocateTableWait = new stepfunctions.Wait(this, 'Wait for query', {
      time: stepfunctions.WaitTime.duration(
        cdk.Duration.minutes(RELOCATE_TABLE_WAIT_TIME),
      ),
    });

    const relocateTableSucceed = new stepfunctions.Succeed(
      this,
      'Query succeed',
    );
    const relocateTableFaild = new stepfunctions.Fail(this, 'Query falid');

    const exportDynamoChoice = new stepfunctions.Choice(
      this,
      'Export complete?',
    )
      .when(
        stepfunctions.Condition.stringEquals('$.exportStaus', 'COMPLETED'), // Choice().next()は不可なので、ネストしている
        relocateTableTask
          .next(relocateTableWait)
          .next(getRelocateTableTask)
          .next(
            new stepfunctions.Choice(this, 'Query complete?')
              .when(
                stepfunctions.Condition.stringEquals('$.status', 'FAILED'),
                relocateTableFaild,
              )
              .when(
                stepfunctions.Condition.stringEquals('$.status', 'SUCCEEDED'),
                relocateTableSucceed,
              )
              .otherwise(relocateTableWait),
          ),
      )
      .otherwise(exportDynamoWait);

    const definition = exportDynamoTask
      .next(exportDynamoWait)
      .next(getExportDynamoStatusTask)
      .next(exportDynamoChoice);

    const stateMachine = new stepfunctions.StateMachine(
      this,
      'exportDynamostateMachine',
      {
        stateMachineName: `export-dynamo`,
        definition: definition,
      },
    );

    /* EventBridgeで定期的に更新する場合(cronの方がよいかも)
    new events.Rule(this, `articleExportDyanmoRule`, {
      ruleName: `article-export-dynamo-rule`,
      schedule: events.Schedule.rate(cdk.Duration.hours(8)),
      targets: [
        new eventsTargets.SfnStateMachine(stateMachine, {
          input: events.RuleTargetInput.fromObject({
            tableName: 'Article',
            s3Bucket: anlysisBucket.bucketName,
          }),
        }),
      ],
    });
    */

    // DynamoDBから出力されたデータをS3に格納し、それを参照するGlueのテーブル定義
    const db = new glue.Database(this, `ExportDynamo`, {
      databaseName: DB_NAME,
    });

    new glue.Table(this, `article`, {
      database: db,
      tableName: `article`,
      bucket: anlysisBucket,
      dataFormat: glue.DataFormat.JSON,
      columns: [
        {
          name: 'Item',
          type: glue.Schema.struct([
            {
              name: 'userId',
              type: glue.Schema.struct([
                {
                  name: 'S',
                  type: glue.Schema.STRING,
                },
              ]),
            },
            {
              name: 'articleId',
              type: glue.Schema.struct([
                {
                  name: 'S',
                  type: glue.Schema.STRING,
                },
              ]),
            },
            {
              name: 'content',
              type: glue.Schema.struct([
                {
                  name: 'S',
                  type: glue.Schema.STRING,
                },
              ]),
            },
            {
              name: 'createAt',
              type: glue.Schema.struct([
                {
                  name: 'N',
                  type: glue.Schema.STRING,
                },
              ]),
            },
          ]),
        },
      ],
    });
  }
}

import * as DynamoDB from '../../infrastructures/dynamodb/dynamodb';

interface Event {
  tableName: string;
}

interface Context {
  invokedFunctionArn: string;
}

interface Result {
  exportArn: string;
  s3Bucket: string;
  tableName: string;
}

export const handler = async (
  event: Event,
  context: Context,
): Promise<Result> => {
  const AWS_REGION = process.env.AWS_REGION!;
  const EXPORT_S3_BUCKET_NAME = process.env.EXPORT_S3_BUCKET_NAME!; // CDKで作成したS3を指定

  const accountId = context.invokedFunctionArn.split(':')[4];
  const tableArn = `arn:aws:dynamodb:${AWS_REGION}:${accountId}:table/${event.tableName}`;

  const exportRequestResult = await DynamoDB.requestExportToS3({ // 次のコードブロックに記述
    tableArn: tableArn,
    s3Bucket: EXPORT_S3_BUCKET_NAME,
    s3Prefix: event.tableName, // s3://${Prefix}/AWSDynamoDB/${exportArn}/data でエクスポートするように指示
  });

  if (!exportRequestResult.ExportDescription?.ExportArn) {
    throw new Error();
  }

  return {
    exportArn: exportRequestResult.ExportDescription.ExportArn,
    s3Bucket: EXPORT_S3_BUCKET_NAME,
    tableName: event.tableName,
  };
};

import * as DynamoDB from 'aws-sdk/clients/dynamodb';

const DYNAMODB_API_VERSION = '2012-10-08';
const REGION = 'ap-northeast-1';

export const dynamodbClient = new DynamoDB({
  apiVersion: DYNAMODB_API_VERSION,
  region: REGION,
});

interface ExportToS3Settings {
  tableArn: string;
  s3Bucket: string;
  s3Prefix?: string;
}

export const requestExportToS3 = async (
  exportToS3Settings: ExportToS3Settings,
): Promise<DynamoDB.ExportTableToPointInTimeOutput> => {
  const response = await dynamodbClient
    .exportTableToPointInTime({
      TableArn: exportToS3Settings.tableArn,
      S3Bucket: exportToS3Settings.s3Bucket,
      S3Prefix: exportToS3Settings.s3Prefix,
      ExportFormat: 'DYNAMODB_JSON',
    })
    .promise();

  return response;
};

import * as DynamoDB from '../../infrastructures/dynamodb/dynamodb';

interface Event {
  exportArn: string;
  s3Bucket: string;
  tableName: string;
}

interface Result {
  exportArn: string;
  exportStaus: string;
  s3Bucket: string;
  tableName: string;
}

export const handler = async (event: Event): Promise<Result> => {
  const exportRequests = await DynamoDB.getExportToS3Status(); // 次のコードブロックに記述
  const filteredExportSummary = exportRequests.ExportSummaries?.filter(
    (exportSummary) => {
      return exportSummary.ExportArn === event.exportArn;
    },
  );

  if (
    !filteredExportSummary ||
    !filteredExportSummary[0].ExportArn ||
    !filteredExportSummary[0].ExportStatus
  ) {
    throw new Error();
  }

  return {
    exportArn: event.exportArn,
    exportStaus: filteredExportSummary[0].ExportStatus,
    s3Bucket: event.s3Bucket,
    tableName: event.tableName,
  };
};

import * as DynamoDB from 'aws-sdk/clients/dynamodb';

const DYNAMODB_API_VERSION = '2012-10-08';
const REGION = 'ap-northeast-1';

export const dynamodbClient = new DynamoDB({
  apiVersion: DYNAMODB_API_VERSION,
  region: REGION,
});

export const getExportToS3Status = async (): Promise<DynamoDB.ListExportsOutput> => {
  const exportsRequest = await dynamodbClient.listExports().promise();

  return exportsRequest;
};

import * as Athena from 'aws-sdk/clients/athena';

interface Event {
  exportArn: string;
  s3Bucket: string;
  tableName: string;
}

interface Result {
  queryExecutionId: string;
}

const DB_NAME = process.env.DB_NAME!;

const API_VERSION = '2017-05-18';
const REGION = process.env.AWS_REGION!;
const ATHENA_WORK_GROUP = process.env.ATHENA_WORK_GROUP!;

export const athenaClient = new Athena({
  apiVersion: API_VERSION,
  region: REGION,
});

export const handler = async (event: Event): Promise<Result> => {
  const splitExportArn = event.exportArn.split('/');

  // CDKでテーブル名はLowerCaseにしているので変換
  const athenaTableName = event.tableName.toLowerCase();
  const exportId = splitExportArn[3];

  const queryLocation = `ALTER TABLE ${DB_NAME}.${athenaTableName}
T LOCATION 's3://${event.s3Bucket}/${event.tableName}/AWSDynamoDB/${exportId}/data/';`;

  const response = await athenaClient
    .startQueryExecution({
      QueryString: queryLocation,
      QueryExecutionContext: {
        Database: DB_NAME,
      },
      WorkGroup: ATHENA_WORK_GROUP, // CDKで作成したAtehnaのワークグループ
    })
    .promise();

  return {
    queryExecutionId: response.QueryExecutionId!, // null undefined判定したほうがよいです..
  };
};

import * as Athena from 'aws-sdk/clients/athena';

interface Event {
  queryExecutionId: string;
}

interface Result {
  status: string;
}

const API_VERSION = '2017-05-18';
const REGION = process.env.AWS_REGION!;

export const athenaClient = new Athena({
  apiVersion: API_VERSION,
  region: REGION,
});

export const handler = async (event: Event): Promise<Result> => {
  const response = await athenaClient
    .getQueryExecution({
      QueryExecutionId: event.queryExecutionId,
    })
    .promise();

  if (response.QueryExecution?.Status?.State !== 'SUCCEEDED') {
    return {
      status: response.QueryExecution?.Status?.State ?? 'UNKNOWN',
    };
  }

  return {
    status: response.QueryExecution?.Status.State,
  };
};